package com.example.demo.rag;

import com.example.demo.service.Assistant;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;

@RestController
@RequiredArgsConstructor
@RequestMapping("/rag/api")
public class RagAPI {
    final Assistant assistant;

    @GetMapping("/high/chat")
    public String chat(@RequestParam(value = "message")String message){
        String aiResponse=assistant.chat(message);
        System.out.println(aiResponse);
        byte[] bytes = aiResponse.getBytes(StandardCharsets.UTF_8);
        System.out.println("Bytes: " + Arrays.toString(bytes));
        return aiResponse;
    }

    final EmbeddingStore<TextSegment> embeddingStore;

    final EmbeddingModel embeddingModel;
    @GetMapping("/load")
    public String load(){
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:\\pools\\demo\\document");
        //String docUrl="https://registration-system1.oss-cn-guangzhou.aliyuncs.com/txt.docx";
        //EmbeddingStoreIngestor.ingest(documents,embeddingStore);
        EmbeddingStoreIngestor.builder()
                //.documentTransformer(document -> Document.document(document.text(),new Metadata().put("userId",1L)))
                //配置文本向量存储的位置
                .embeddingStore(embeddingStore)
                //配置文本向量模型
                .embeddingModel(embeddingModel)
                //怎么切分文本
                .documentSplitter(new DocumentByLineSplitter(50,30))
                .build().ingest(documents);//传入文本
        return "success";
    }
}
